import pysam
import argparse
import pod5 as p5
from tqdm import tqdm
from pathlib import Path
from collections import defaultdict

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--inbam', type=str, required=True, help="Input BAM file")
    parser.add_argument('--outbam', type=str, required=True, help="Output BAM file")
    parser.add_argument('--pod5_dir', type=str, required=True, help="Directory containing POD5 files")
    return parser.parse_args()

def get_readid_to_pod5_mapping(pod5_dir):
    """创建read_id到pod5文件名的映射"""
    readid_to_pod5 = defaultdict(list)
    recursive = True
    glob = Path.rglob if recursive else Path.glob
    
    for pod5_path in tqdm(glob(Path(pod5_dir), "*.pod5"), desc="Scanning POD5 files"):
        with p5.DatasetReader(pod5_path, recursive=True, max_cached_readers=1) as dataset:
            for read_record in dataset:
                read_id = str(read_record.read_id)
                pod5_filename = pod5_path.name
                readid_to_pod5[read_id].append(pod5_filename)
    
    final_mapping = {k: v[0] for k, v in readid_to_pod5.items()}
    print(f"Found {len(final_mapping)} unique read IDs in POD5 files")
    return final_mapping

def update_bam_with_pod5_filenames(inbam, outbam, readid_to_pod5):
    """更新BAM文件中的fn标签，考虑query_name和pi标签"""
    input_bam = pysam.AlignmentFile(inbam, 'rb', check_sq=False)
    output_bam = pysam.AlignmentFile(outbam, 'wb', header=input_bam.header)
    
    modified_count = 0
    total_count = 0
    pi_match_count = 0
    
    for read in tqdm(input_bam, desc="Processing BAM reads"):
        total_count += 1
        read_id = read.query_name
        pi_tag = read.get_tag("pi") if read.has_tag("pi") else None
        
        # 检查query_name或pi标签是否在映射中
        target_id = None
        if read_id in readid_to_pod5:
            target_id = read_id
        elif pi_tag and pi_tag in readid_to_pod5:
            target_id = pi_tag
            pi_match_count += 1
            
        if target_id:
            # 设置或更新fn标签为对应的pod5文件名
            read.set_tag('fn', readid_to_pod5[target_id], value_type='Z')
            modified_count += 1
            
        output_bam.write(read)
    
    print(f"Processed {total_count} reads")
    print(f"Modified {modified_count} reads with POD5 filenames")
    print(f"Matched {pi_match_count} reads via pi tag")
    input_bam.close()
    output_bam.close()

def main():
    args = parse_args()
    print(f"Processing POD5 files from: {args.pod5_dir}")
    print(f"Input BAM: {args.inbam}")
    print(f"Output BAM: {args.outbam}")
    
    # 获取read_id到pod5文件名的映射
    readid_to_pod5 = get_readid_to_pod5_mapping(args.pod5_dir)
    
    # 更新BAM文件
    print("Starting BAM file processing")
    update_bam_with_pod5_filenames(args.inbam, args.outbam, readid_to_pod5)
    print("Processing completed")

if __name__ == '__main__':
    main()